import cv2
import numpy as np


class CameraAwareDetector:
    def __init__(self, camera_params):
        # 初始化相机参数
        self.camera_matrix = np.array(camera_params['camera_matrix']).reshape(3, 3)
        self.dist_coeffs = np.array(camera_params['distortion_coefficients'])
        self.image_size = (camera_params['image_width'], camera_params['image_height'])

        # 预计算畸变校正映射（提高实时性）
        self.map1, self.map2 = cv2.initUndistortRectifyMap(
            self.camera_matrix, self.dist_coeffs, None, self.camera_matrix,
            self.image_size, cv2.CV_16SC2
        )

    def preprocess_image(self, image):
        """使用相机参数进行图像预处理"""
        # 1. 畸变校正
        undistorted = cv2.remap(image, self.map1, self.map2, cv2.INTER_LINEAR)

        # 2. 可选：根据相机内参进行图像增强
        # 这里可以利用焦距等信息调整处理参数
        fx = self.camera_matrix[0, 0]  # x轴焦距
        # 焦距越大，图像细节越重要，可以调整处理参数

        return undistorted

    def detect_court_lines(self, image):
        """识别场地线（结合相机参数）"""
        # 预处理图像
        processed_img = self.preprocess_image(image)

        # 转换为灰度图
        gray = cv2.cvtColor(processed_img, cv2.COLOR_BGR2GRAY)

        # 使用Canny边缘检测（参数可以根据相机焦距调整）
        edges = cv2.Canny(gray, 50, 150)

        # 霍夫线变换检测直线
        lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=100,
                                minLineLength=50, maxLineGap=10)

        return lines, processed_img

    def calculate_real_world_coordinates(self, pixel_points, assumed_height=1.0):
        """将像素坐标转换为真实世界坐标（假设场地平面）"""
        # 假设场地在z=0平面上，计算3D坐标
        inv_camera_matrix = np.linalg.inv(self.camera_matrix)

        real_world_points = []
        for pixel in pixel_points:
            # 像素坐标转齐次坐标
            pixel_homogeneous = np.array([pixel[0], pixel[1], 1.0])

            # 计算射线方向
            ray_dir = inv_camera_matrix @ pixel_homogeneous

            # 假设场地高度为assumed_height，计算与平面的交点
            # 这里简化处理，实际需要知道相机高度和姿态
            scale = assumed_height / ray_dir[1]  # 假设y轴垂直
            point_3d = ray_dir * scale

            real_world_points.append(point_3d)

        return np.array(real_world_points)

    def draw_detection_results(self, image, lines, key_points=None):
        """绘制检测结果"""
        result = image.copy()

        # 绘制检测到的线
        if lines is not None:
            for line in lines:
                x1, y1, x2, y2 = line[0]
                cv2.line(result, (x1, y1), (x2, y2), (0, 255, 0), 2)

        # 绘制关键点
        if key_points is not None:
            for point in key_points:
                cv2.circle(result, tuple(point.astype(int)), 5, (0, 0, 255), -1)

        return result


# 相机参数配置
camera_params = {
    'image_width': 1440,
    'image_height': 1080,
    'camera_matrix': [2709.520872175487, 0.0, 650.0710018353218,
                      0.0, 2575.659900756491, 600.5234864639328,
                      0.0, 0.0, 1.0],
    'distortion_coefficients': [3.765324073159127, -190.56351270576693,
                                0.07450270179341688, -0.05268220024905785,
                                3291.737231347516]
}


def main():
    # 初始化检测器
    detector = CameraAwareDetector(camera_params)

    # 读取图像
    image = cv2.imread('badminton_court.jpg')
    if image is None:
        print("无法读取图像文件")
        return

    # 调整图像尺寸匹配相机参数
    image = cv2.resize(image, (1440, 1080))

    # 进行场地线检测
    lines, processed_img = detector.detect_court_lines(image)

    # 提取关键点（示例：场地角点）
    key_points = np.array([[100, 200], [1300, 200], [100, 800], [1300, 800]])

    # 计算真实世界坐标
    real_world_points = detector.calculate_real_world_coordinates(key_points)
    print("真实世界坐标:", real_world_points)

    # 绘制结果
    result_img = detector.draw_detection_results(processed_img, lines, key_points)

    # 显示结果
    cv2.imshow('Original', image)
    cv2.imshow('Undistorted', processed_img)
    cv2.imshow('Detection Result', result_img)

    cv2.waitKey(0)
    cv2.destroyAllWindows()

    # 保存结果
    cv2.imwrite('detection_result.jpg', result_img)


if __name__ == "__main__":
    main()